import akshare as ak
import pandas as pd
import os
from tqdm import tqdm
import numpy as np 
import datetime

# 399311  2000 399303
def get_stock_codes(root=None, name='国证1000') -> None:
    today = datetime.datetime.today()
    str_y_m = f'{today.year}{today.month-1}'
    dic = {'国证1000':'399311', '国证2000':'399303'}
    if root is None:
        root = r"data/stock_data/code_date"
    if not os.path.exists(root):
        os.makedirs(root)
    fund_etf_fund_daily_em_df = ak.index_detail_cni(symbol= dic[name], date=str_y_m)
    fund_etf_fund_daily_em_df.to_csv( os.path.join(root, f"{name}.csv"), encoding="utf-8", index=False )
    return fund_etf_fund_daily_em_df

def get_price(code_list, root=None, start_date='20010101'):
    if root is None:
        root = r'data/stock_data/daily/price'
    if not os.path.exists(root):
        os.makedirs(root)
    for i in tqdm(code_list):
        try:
            result = ak.stock_zh_a_hist(symbol=i, period="daily", start_date=start_date, adjust="qfq")
            result.index = pd.to_datetime(result['日期'])
            result = result.sort_index()
            result.drop(columns='日期', inplace=True)
            result = result.astype(float)
            result = result[['开盘','收盘', '最高', '最低', '成交量', '成交额']]
            result.columns = ['open', 'close', 'high','low','volume','amount'] 
            close_t = result['close'][-1]  
            if (close_t < 5) or (close_t>150):
                continue
            else:
                result.to_pickle(f'{root}/{i}.pkl.gzip')
        except Exception:
            print('error:', i)
            continue

def get_market_value(code_list, root=None):
    if root is None:
        root = r'data/stock_data/daily/market_value'
    if not os.path.exists(root):
        os.makedirs(root)
    for i in tqdm(code_list): #对齐索引即可
        try:
            result = ak.stock_a_lg_indicator(symbol=i)
            result.index = pd.to_datetime(result['trade_date'])
            result = result.sort_index()
            result.drop(columns='trade_date', inplace=True)
            result = result.astype(float)
            result.to_pickle(f'{root}/{i}.pkl.gzip')
        except Exception:
            print('error:', i)
            continue

def get_account(code_list, root=None):
    if root is None:
        root = r'data/stock_data/daily/account'
    if not os.path.exists(root):
        os.makedirs(root)
    for i in tqdm(code_list): #需进行重采样处理
        try:
            result = ak.stock_financial_analysis_indicator(symbol=i)
            result.index = pd.to_datetime(result['日期'])
            result = result.sort_index()
            result.drop(columns='日期', inplace=True)
            result = result.replace('--',np.nan)
            result = result[result.index>'2001-01-01']
            # result = result.astype(float)
            result.to_pickle(f'{root}/{i}.pkl.gzip')
        except Exception:
            print('error:', i)
            continue

def get_other_data(root=None):
    if root is None:
        root = r'data/stock_data/daily/other'
    if not os.path.exists(root):
        os.makedirs(root)
    #legu数据过短，2020开始
    #可自行增加截面因子
    # result = ak.stock_a_high_low_statistics(symbol="zz500")
    # result.index = pd.to_datetime(result['date'])
    # result = result.sort_index()
    # result.drop(columns='date', inplace=True)
    # result = result.astype(float)/500
    # result.to_pickle(f'{root}/high_low_zz500.pkl.gzip')

    # result = ak.stock_a_high_low_statistics(symbol="sz50")
    # result.index = pd.to_datetime(result['date'])
    # result = result.sort_index()
    # result.drop(columns='date', inplace=True)
    # result = result.astype(float)/50
    # result.to_pickle(f'{root}/high_low_sz50.pkl.gzip')

    # result = ak.stock_a_high_low_statistics(symbol="hs300")
    # result.index = pd.to_datetime(result['date'])
    # result = result.sort_index()
    # result.drop(columns='date', inplace=True)
    # result = result.astype(float)/300
    # result.to_pickle(f'{root}/high_low_hs300.pkl.gzip')

    # result = ak.stock_a_high_low_statistics(symbol="all")
    # result.index = pd.to_datetime(result['date'])
    # result = result.sort_index()
    # result.drop(columns='date', inplace=True)
    # result = result.astype(float)
    # result.to_pickle(f'{root}/high_low_all.pkl.gzip')

    #2017开始，勉强可用
    lst = ['sh','sz','cy','kc','000300.SH','000016.SH','000905.SH','000852.SH']
    for i in lst:
        result = ak.stock_a_pe_and_pb(symbol=i)
        result.index = pd.to_datetime(result['date'])
        result.drop(columns='date', inplace=True)
        result = result.astype(float)
        result.to_pickle(f'{root}/pepb_{i[:6]}.pkl.gzip')

    industry = ak.stock_board_industry_summary_ths()
    industry.to_pickle(f'{root}/indyustry.pkl.gzip')

    lst = industry['板块']
    industry_root = os.path.join(f'{root}/indyustry')
    if not os.path.exists(industry_root):
        os.makedirs(industry_root)
    for i in lst:
        result = ak.stock_board_industry_cons_ths(symbol=i)
        result.to_pickle(f'{industry_root}/{i}.pkl.gzip')
    #可以使用sin编码

    bound_ratio = ak.bond_zh_us_rate()
    bound_ratio.index = pd.to_datetime(bound_ratio['日期'])
    bound_ratio.drop(columns='日期', inplace=True)
    bound_ratio = bound_ratio.astype(float)
    bound_ratio.to_pickle(f'{root}/bound_ratio.pkl.gzip')

    currency_boc_safe_df = ak.currency_boc_safe()
    currency_boc_safe_df.index = pd.to_datetime(currency_boc_safe_df['日期'])
    currency_boc_safe_df.drop(columns='日期', inplace=True)
    currency_boc_safe_df = currency_boc_safe_df.astype(float)
    currency_boc_safe_df.to_pickle(f'{root}/currency_boc_safe.pkl.gzip')

    rate_interbank_df = ak.rate_interbank(market="上海银行同业拆借市场", symbol="Shibor人民币", indicator="3月")
    rate_interbank_df.index = pd.to_datetime(rate_interbank_df['报告日'])
    rate_interbank_df.drop(columns='报告日', inplace=True)
    rate_interbank_df = rate_interbank_df.astype(float)
    rate_interbank_df.to_pickle(f'{root}/rate_interbank.pkl.gzip')

    stock_ebs_lg_df = ak.stock_ebs_lg()
    stock_ebs_lg_df.index = pd.to_datetime(stock_ebs_lg_df['日期'])
    stock_ebs_lg_df.drop(columns='日期', inplace=True)
    stock_ebs_lg_df = stock_ebs_lg_df.astype(float)
    stock_ebs_lg_df.to_pickle(f'{root}/ebs_lg.pkl.gzip')


